/*
 * Video Acceleration API (video encoding) encode sample
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
 * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */

/**
 * @file
 * Intel VAAPI-accelerated encoding example.
 *
 * @example vaapi_encode.c
 * This example shows how to do VAAPI-accelerated encoding. now only support NV12
 * raw file, usage like: vaapi_encode 1920 1080 input.yuv output.h264
 *
 */

#include <stdio.h>
#include <string.h>
#include <errno.h>
#include <math.h>

#include <libavcodec/avcodec.h>
#include <libavutil/pixdesc.h>
#include <libavutil/hwcontext.h>

#include <sys/time.h>
#include "utils.h"


struct EncoderContext
{
    astring codec_name;
    astring enc_video;
    astring dec_dir;

    uint8_t *raw_data[32];
    long     raw_data_size[32];
    long     raw_data_readpos;

    int      samples;
    int      width;
    int      height;
    int      pixfmt;

    int      cur_frame;
    int     *cur_bytes;
    double  *cur_psnr;
    double  *cur_elpase;
    int      bytes_rec[52];
    double   psnr_rec[52];
    double   elapse_rec[52];

    int      total_frames;
    double   total_time_used;

    AVDictionary *options;
};

extern int swdecode(astring filename, astring outdir, int dump_file_count);

int read_rawdata_file(struct EncoderContext *enc, astring raw_image)
{
    char fn[_MAX_PATH];
    strcpy(fn, raw_image);
    char *ic = strrchr(fn, '.');
    if(ic == NULL) ic = fn + strlen(fn);
    ic--;

    for(int n = 0; n < 32; n++)
    {
        FILE *fi = fopen(fn, "rb");
        if(!fi)
        {
            if(n) return 0;
            fprintf(stderr, "Fail to open input file: %s, %s\n", raw_image, strerror(errno));
            return -1;
        }

        fseek(fi, 0, SEEK_END);
        long len = ftell(fi);
        fseek(fi, 0, SEEK_SET);

        enc->raw_data[n] = (uint8_t *)malloc(len);
        fread(enc->raw_data[n], len, 1, fi);
        enc->raw_data_size[n] = len;
        fclose(fi);

        if((*ic >= '0' && *ic < '9') || (*ic >= 'a' && *ic < 'z'))
            *ic = *ic + 1;
    }
    return 0;
}

int get_rawdata(struct EncoderContext *enc, uint8_t *dstbuf, int size)
{
    long pos = enc->raw_data_readpos;

    int idx = 0;
    while(pos > enc->raw_data_size[idx])
    {
        if(enc->raw_data_size[idx] == 0)
            return -1; // no more data

        pos -= enc->raw_data_size[idx++];
    }

    while(size)
    {
        int cplen = MIN(enc->raw_data_size[idx] - pos, size);
        memcpy(dstbuf, enc->raw_data[idx] + pos, cplen);
        size -= cplen;
        enc->raw_data_readpos += cplen;

        if(enc->raw_data_size[++idx] == 0) break;
        pos = 0;
    }

    return size == 0;
}

static int set_hwframe_ctx(AVCodecContext *ctx, AVBufferRef *hw_device_ctx, int width, int height)
{
    AVBufferRef *hw_frames_ref;
    AVHWFramesContext *frames_ctx = NULL;
    int err = 0;

    if (!(hw_frames_ref = av_hwframe_ctx_alloc(hw_device_ctx))) {
        fprintf(stderr, "Failed to create VAAPI frame context.\n");
        return -1;
    }
    frames_ctx = (AVHWFramesContext *)(hw_frames_ref->data);
    frames_ctx->format    = AV_PIX_FMT_VAAPI;
    frames_ctx->sw_format = AV_PIX_FMT_NV12;
    frames_ctx->width     = width;
    frames_ctx->height    = height;
    frames_ctx->initial_pool_size = 20;
    if ((err = av_hwframe_ctx_init(hw_frames_ref)) < 0) {
        fprintf(stderr, "Failed to initialize VAAPI frame context."
                "Error code: %s\n",av_err2str(err));
        av_buffer_unref(&hw_frames_ref);
        return err;
    }
    ctx->hw_frames_ctx = av_buffer_ref(hw_frames_ref);
    if (!ctx->hw_frames_ctx)
        err = AVERROR(ENOMEM);

    av_buffer_unref(&hw_frames_ref);
    return err;
}

static int encode_write(struct EncoderContext *enc, AVCodecContext *avctx, AVFrame *frame, FILE *fout)
{
    int ret = 0;
    AVPacket *enc_pkt;

    if (!(enc_pkt = av_packet_alloc()))
        return AVERROR(ENOMEM);

    if ((ret = avcodec_send_frame(avctx, frame)) < 0)
    {
        fprintf(stderr, "Error code: %s\n", av_err2str(ret));
        goto end;
    }

    while (1)
    {
        ret = avcodec_receive_packet(avctx, enc_pkt);
        if (ret) break;

        enc_pkt->stream_index = 0;
        ret = fwrite(enc_pkt->data, enc_pkt->size, 1, fout);

        *enc->cur_bytes += enc_pkt->size;
        enc->cur_frame++;
        enc->total_frames++;
        printf("encode frame_%d, size: %.2fkb\n", enc->cur_frame, *enc->cur_bytes/1024.0);

        av_packet_unref(enc_pkt);
    }

end:
    av_packet_free(&enc_pkt);

    if (ret == AVERROR_EOF || ret == AVERROR(EAGAIN))
        ret = 0;
    return ret;
}

int vaapi_encode(struct EncoderContext *enc)
{
    int err;
    FILE *fo = NULL;

    AVFrame *sw_frame = NULL;
    AVFrame *hw_frame = NULL;
    AVBufferRef *hw_device_ctx = NULL;

    const AVCodec *codec = NULL;
    AVCodecContext *avctx = NULL;

    int width = enc->width;
    int height = enc->height;

    if (!(fo = fopen(enc->enc_video, "w+b"))) {
        fprintf(stderr, "Fail to open output file: %s, %s\n", enc->enc_video, strerror(errno));
        err = -1;
        goto finalize;
    }

    err = av_hwdevice_ctx_create(&hw_device_ctx, AV_HWDEVICE_TYPE_VAAPI, NULL, enc->options, 0);
    if (err < 0) {
        fprintf(stderr, "Failed to create a VAAPI device. Error code: %s\n", av_err2str(err));
        fprintf(stderr, "Use '-d' option to select drm and take a try again.\n");
        goto finalize;
    }

    if (!(codec = avcodec_find_encoder_by_name(enc->codec_name))) {
        fprintf(stderr, "Could not find encoder (%s).\n", enc->codec_name);
        err = -1;
        goto finalize;
    }

    if (!(avctx = avcodec_alloc_context3(codec))) {
        err = AVERROR(ENOMEM);
        goto finalize;
    }

    avctx->width     = width;
    avctx->height    = height;
    avctx->time_base = (AVRational){1, 25};
    avctx->framerate = (AVRational){25, 1};
    avctx->sample_aspect_ratio = (AVRational){1, 1};
    avctx->pix_fmt   = AV_PIX_FMT_VAAPI;

    /* set hw_frames_ctx for encoder's AVCodecContext */
    if ((err = set_hwframe_ctx(avctx, hw_device_ctx, width, height)) < 0) {
        fprintf(stderr, "Failed to set hwframe context.\n");
        goto finalize;
    }

    if ((err = avcodec_open2(avctx, codec, &enc->options)) < 0) {
        fprintf(stderr, "Cannot open video encoder codec. Error code: %s\n", av_err2str(err));
        goto finalize;
    }

    struct timeval time_start;
    struct timeval time_end;

    gettimeofday(&time_start, NULL);
    int size = width * height;
    while (1)
    {
        if (!(sw_frame = av_frame_alloc())) {
            err = AVERROR(ENOMEM);
            goto finalize;
        }
        /* read data into software frame, and transfer them into hw frame */
        sw_frame->width  = width;
        sw_frame->height = height;
        sw_frame->format = enc->pixfmt;

        if ((err = av_frame_get_buffer(sw_frame, 0)) < 0)
            goto finalize;

        switch(enc->pixfmt)
        {
            case AV_PIX_FMT_YUYV422:
                err = get_rawdata(enc, sw_frame->data[0], size * 2);
                //err = fread(sw_frame->data[0], size * 2, 1, fi);
                break;

            case AV_PIX_FMT_NV12:
                //err = fread(sw_frame->data[0], size, 1, fi);
                //err = fread(sw_frame->data[1], size/2, 1, fi);
                err = get_rawdata(enc, sw_frame->data[0], size);
                err = get_rawdata(enc, sw_frame->data[1], size/2);
                break;

            case AV_PIX_FMT_YUV420P:
                //err = fread(sw_frame->data[0], size, 1, fi);
                //err = fread(sw_frame->data[1], size/4, 1, fi);
                //err = fread(sw_frame->data[2], size/4, 1, fi);
                err = get_rawdata(enc, sw_frame->data[0], size);
                err = get_rawdata(enc, sw_frame->data[1], size/4);
                err = get_rawdata(enc, sw_frame->data[2], size/4);
                break;

            default:
                err = -1;
                fprintf(stderr, "Invalid input pixfmt %s\n", av_get_pix_fmt_name(enc->pixfmt));
        }
        if(err <= 0) break;


        if (!(hw_frame = av_frame_alloc())) {
            err = AVERROR(ENOMEM);
            goto finalize;
        }
        if ((err = av_hwframe_get_buffer(avctx->hw_frames_ctx, hw_frame, 0)) < 0) {
            fprintf(stderr, "Error code: %s.\n", av_err2str(err));
            goto finalize;
        }

        if (!hw_frame->hw_frames_ctx) {
            err = AVERROR(ENOMEM);
            goto finalize;
        }
        if ((err = av_hwframe_transfer_data(hw_frame, sw_frame, 0)) < 0) {
            fprintf(stderr, "Error while transferring frame data to surface."
                    "Error code: %s.\n", av_err2str(err));
            goto finalize;
        }

        if ((err = (encode_write(enc, avctx, hw_frame, fo))) < 0) {
            fprintf(stderr, "Failed to encode.\n");
            goto finalize;
        }
        av_frame_free(&hw_frame);
        av_frame_free(&sw_frame);
    }

    /* flush encoder */
    err = encode_write(enc, avctx, NULL, fo);

    gettimeofday(&time_end, NULL);
    double elapse = (time_end.tv_sec - time_start.tv_sec) * 1000 + (double)(time_end.tv_usec - time_start.tv_usec) / 1000;
    *enc->cur_elpase = elapse;
    enc->total_time_used += elapse;
    printf("encoding time %.2fms\n", elapse);

finalize:
    if(fo) fclose(fo);

    av_frame_free(&sw_frame);
    av_frame_free(&hw_frame);
    avcodec_free_context(&avctx);
    av_buffer_unref(&hw_device_ctx);
    hw_device_ctx = NULL;
    fprintf(stderr, "%s err %d\n", __func__, err);
    return err;
}

int image_compare(struct EncoderContext *enc)
{
    FILE *fi = NULL;

    char short_name[256] = {0};

    {
        char *p;
        char *last_dot;

#ifdef _WIN32
        strcpy(short_name, (p = strrchr(enc->enc_video, '\\')) ? p + 1 : enc->enc_video);
#else
        strcpy(short_name, (p = strrchr(enc->enc_video, '/')) ? p + 1 : enc->enc_video);
#endif
        last_dot = strrchr(short_name, '.');
        if (last_dot != NULL) *last_dot = '\0';
        printf("short_name: %s\n", short_name);
    }

    int width = enc->width;
    int height = enc->height;

    int pixels = width * height;
    int raw_image_len = (enc->pixfmt == AV_PIX_FMT_YUYV422) ? pixels * 2 : pixels * 3 / 2;
    int dec_image_len = pixels * 3 / 2;

    uint8_t *rawbuf = (uint8_t *)malloc(raw_image_len);
    uint8_t *decbuf = (uint8_t *)malloc(dec_image_len);


    int err = 0;
    enc->raw_data_readpos = 0;
    for(int frame = 0; frame < enc->cur_frame; frame++)
    {
        //if((err = fread(rawbuf, raw_image_len, 1, fi)) <= 0)
        //    break;

        if((err = get_rawdata(enc, rawbuf, raw_image_len)) <=0)
            break;

        char fn[1024] = {0};
        sprintf(fn, "%s/%s-%d.nv12", enc->dec_dir, short_name, frame + 1);

        if (!(fi = fopen(fn, "rb")))
        {
            fprintf(stderr, "Fail to open swout_filepath file : %s-%s\n", fn, strerror(errno));
            return -3.0;
        }
        err = fread(decbuf, dec_image_len, 1, fi);
        fclose(fi);

        if(err <= 0)
            break;

        float mse = 0;
        uint8_t
            *ps1 = rawbuf,
            *ps2 = decbuf;

        for(int pix = 0; pix < pixels; pix++)
        {
            float diff = (float)(*ps1++ - *ps2++);
            if(enc->pixfmt == AV_PIX_FMT_YUYV422)
                ps1++;
            mse += sq(diff);
        }

        mse /= pixels;
        float psnr = 10 * log10(sq(255) / mse);
        printf("frame_%d psnr %.2f\n", frame, psnr);
        *enc->cur_psnr += psnr;
    }

    if(rawbuf) free(rawbuf);
    if(decbuf) free(decbuf);

    return 0;
}

float calc_score(struct EncoderContext *enc)
{
    double rate, distort, tm, score, overall_score = 0;
    for(int n = enc->samples - 1; n >= 0; n--)
    {
        rate = tanh(log(enc->bytes_rec[n]) - 4) + 3;
        distort = enc->psnr_rec[n];
        tm = log(enc->elapse_rec[n]) + 6;
        score = tanh(distort / rate / tm) * 100;
        overall_score += score;
        //printf("rate %.2f, rd %.2f, tm %.2f, score %.2f\n", rate, distort / rate, tm, score);
    }
    overall_score /= enc->samples;
    //printf("overall score: %.2f\n", overall_score);
    return overall_score;
}



int hwencode(int width, int height, int pixfmt, astring codec, astring raw_image, astring enc_video, astring dec_dir, int drm_mode)
{
    int ret = 0;

    struct EncoderContext enc = {0};
    enc.width = width;
    enc.height = height;
    enc.codec_name = codec;
    enc.pixfmt = pixfmt;
    enc.enc_video = enc_video;
    enc.dec_dir = dec_dir;

    printf("input: %s\n", raw_image);
    printf("output: %s\n", enc_video);
    printf("input pixfmt: %s\n", av_get_pix_fmt_name(pixfmt));
    printf("dimension: %dx%d\n", width, height);
    printf("encoder: %s\n", codec);
    printf("use_drm: %s\n", drm_mode?"true":"false");

    if(read_rawdata_file(&enc, raw_image) < 0)
        return -1;

    av_dict_set(&enc.options, "connection_type", drm_mode ? "drm" : "x11", 0);

    int qp_preset[] = {1, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50};
    for(int idx = 0; idx < sizeof(qp_preset)/sizeof(qp_preset[0]); idx++)
    {
        int qp = qp_preset[idx];
        printf("encoding_%d, qp %d\n", idx, qp);
        av_dict_set_int(&enc.options, "qp", qp, 0);

        enc.cur_frame  = 0;
        enc.cur_bytes  = &enc.bytes_rec[idx];
        enc.cur_psnr   = &enc.psnr_rec[idx];
        enc.cur_elpase = &enc.elapse_rec[idx];
        enc.raw_data_readpos = 0;
        ret = vaapi_encode(&enc);
        if(ret < 0)
        {
            printf("vaapi_encode failed, ret = %d.\n", ret);
            break;
        }

        ret = swdecode(enc_video, dec_dir, enc.cur_frame);
        if(ret < 0)
        {
            printf("run swdecode failed, ret = %d.\n", ret);
            break;
        }

        image_compare(&enc);
        *enc.cur_bytes  /= enc.cur_frame;
        *enc.cur_psnr   /= enc.cur_frame;
        *enc.cur_elpase /= enc.cur_frame;
        enc.samples++;
    }

    if(ret >= 0)
    {
        printf("\n");
        for(int idx = 0; idx < sizeof(qp_preset)/sizeof(qp_preset[0]); idx++)
        {
            int qp = qp_preset[idx];
            printf("encoding_%d, qp %d, psnr %.2f, bytes %.2fkb, time_used %.2fms\n",
                idx, qp, enc.psnr_rec[idx], enc.bytes_rec[idx]/1024.0, enc.elapse_rec[idx]);
        }

        float score = calc_score(&enc);
        printf("=======================\n");
        printf("score for %s encoder: %.2f\n", codec, score);
        printf("COUNT|%.2f|0|score\n", score);
        printf("=======================\n");
    }

    av_dict_free(&enc.options);
    for(int n = 0; n < 32; n++)
    {
        if(enc.raw_data_size[n])
            free(enc.raw_data[n]);
    }
    return ret;
}


